# -*- coding: utf-8 -*-
import numpy as np
from .node import Node
from .ElementBase import elementBase
from .tools import *
from typing import List
from .Loads import *
# 绗架单元类 OK

class T3D2(elementBase):
    """
    三维,2节点线性空间桁架单元,截面均一,各项同性材料,节点自由度:Ux,Uy,Uz.单元公式见项目文档
    
    单元定义需要:A,E,rho,节点坐标;
    
    载荷支持集中节点力,不考虑热效应
    """
    VTK_TYPE=3
    
    def __init__(self,label:int,
                    node1:Node,
                    node2:Node,
                    E=0.0,A=0.0,rho=0.0):
        super().__init__(label,3)
        self.node1=node1
        self.node2=node2
        self.E=E
        self.A=A
        self.rho=rho
    @property
    def L(self):
        """
        计算单元长度
        
        returns:
            float: 单元长度
        """
        return np.linalg.norm(self.node2.coordinate-self.node1.coordinate)
    @property
    def t(self):
        """
        计算转换矩阵[lamda]
        
        returns:
            np.ndarray: 转换矩阵(2,6)
        """
        lij,mij,nij=(self.node2.coordinate-self.node1.coordinate)/self.L
        return np.array([[lij,mij,nij,0,0,0],[0,0,0,lij,mij,nij]])
    @property
    def Ke(self):
        """
        计算单元在全局坐标系下的单元刚度矩阵
        
        returns:
            np.ndarray: 单元刚度矩阵(6,6)
        """
        # 一维杆的刚度矩阵公式
        ke=np.array([[1,-1],[-1,1]])*(self.E*self.A/self.L)
        return (self.t.T)@ke@self.t
    @property
    def Me(self):
        """
        计算全局坐标系下的单元质量矩阵(lumped mass matrix)
        
        returns:
            np.ndarray: 单元质量矩阵(6,6)
        """
        return np.eye(6)*(self.rho*self.A*self.L/2.0)
    @property
    def D(self):
        """
        返回单元的弹性模量E
        
        returns:
            float: 弹性模量E
        """
        return self.E
    @property
    def dofNum(self):
        """
        返回单元的自由度数
        
        returns:
            int: 自由度数
        """
        return 6
    @property
    def NodesList(self):
        """
        返回单元的节点列表
        
        returns:
            List[Node]: 节点列表
        """
        return [self.node1,self.node2]
    @property
    def ElemDofIndexs(self):
        """
        返回单元的自由度索引列表(0-based)
        
        returns:
            List[int]: 自由度索引列表
        """
        return dofInds(self.node1.label,self.NodeDofNum)+dofInds(self.node2.label,self.NodeDofNum)
    def AxialStress(self,ue=None):
        """
        计算单元的轴向应力
        
        args:
            ue (np.ndarray, optional): 单元位移向量(6,). Defaults to None.
        
        returns:
            float: 轴向应力
        """
        if ue is None:
            raise ValueError("ue must be provided to calculate axial stress")
        return self.AxialStrain(ue)*self.E
    def AxialStrain(self,ue=None):
        """
        计算单元的轴向应变
        
        Args:
            ue (np.ndarray, optional): 单元位移向量(6,). Defaults to None.
        
        returns:
            float: 轴向应变
        """
        if ue is None:
            raise ValueError("ue must be provided to calculate axial strain")
        if ue.shape!=(6,1):
            ue=ue.reshape((6,1))
        return (np.array([[-1.0,1.0]])*self.L)@self.t@ue

class T2D2(elementBase):
    """
    2维,2节点线性平面桁架单元,截面均一,各项同性材料,节点自由度:Ux,Uy.单元公式见项目文档
    
    单元定义需要:A,E,rho,节点坐标;
    
    载荷支持集中节点力,不考虑热效应
    """
    VTK_TYPE=3
    
    def __init__(self,label:int,
                    node1:Node,node2:Node,
                    E=0.0,A=0.0,rho=0.0):
        super().__init__(label,2)
        self.node1=node1
        self.node2=node2
        self.E=E
        self.A=A
        self.rho=rho
    @property
    def get_connectivity(self):
        return [self.node1.label,self.node2.label]
    @property
    def L(self):
        """
        计算单元长度
        
        returns:
            float: 单元长度
        """
        return np.linalg.norm(self.node2.coordinate-self.node1.coordinate)
    @property
    def t(self):
        """
        计算转换矩阵[lamda]=[l,m,0,0;0,0,l,m]
        
        returns:
            np.ndarray: 转换矩阵(4,4)
        """
        lij,mij,_=(self.node2.coordinate-self.node1.coordinate)/self.L
        return np.array([[lij,mij,0,0],
                        [0,0,lij,mij]])
    @property
    def Ke(self):
        """
        计算当前单元的全局坐标系下的单元刚度矩阵
        
        returns:
            np.ndarray: 单元刚度矩阵(4,4)
        """
        # 一维杆的刚度矩阵公式
        ke=np.array([[1,-1],[-1,1]])*(self.E*self.A/self.L)
        return self.t.T@ke@self.t
    @property
    def Me(self):
        """
        计算当前单元的集中质量矩阵(总质量均布到每个节点)
        
        returns:
            np.ndarray: 单元质量矩阵(4,4)
        """
        return (self.rho*self.A*self.L/2.0)*np.eye(4) 
    @property
    def D(self):
        """
        返回当前单元的弹性模量E
        
        returns:
            float: 弹性模量E
        """
        return self.E
    @property
    def dofNum(self):
        """
        获取当前单元的自由度数量
        
        returns:
            int: 自由度数量
        """
        return 4
    @property
    def NodesList(self):
        """
        获取当前单元的节点列表
        
        returns:
            List[Node]: 节点列表
        """
        return [self.node1,self.node2]
    @property
    def ElemDofIndexs(self):
        """
        获取当前单元的全局自由度索引列表(0-based)
        
        returns:
            List[int]: 自由度索引列表
        """
        return dofInds(self.node1.label,self.NodeDofNum)+dofInds(self.node2.label,self.NodeDofNum)
    def AxialStress(self,ue=None):
        """
        计算单元的轴向应力
        
        args:
            ue (np.ndarray, optional): 单元位移向量(6,). Defaults to None.
        
        returns:
            float: 轴向应力
        """
        if ue is None:
            raise ValueError("ue must be provided to calculate axial stress")
        return self.AxialStrain(ue)*self.E
    def AxialStrain(self,ue=None):
        """
        计算单元的轴向应变
        
        Args:
            ue (np.ndarray, optional): 单元位移向量(6,). Defaults to None.
        
        returns:
            float or numpy.array(1): 轴向应变
        """
        if ue is None:
            raise ValueError("ue must be provided to calculate axial strain")
        if ue.shape!=(4,1):
            ue=ue.reshape((4,1))
        return (np.array([[-1.0,1.0]])*self.L)@self.t@ue

class T1D2(elementBase):
    """
    1维,2节点线性空间桁架单元,截面均一,各项同性材料,节点自由度:Ux.单元公式见项目文档
    
    单元定义需要:A,E,rho,节点坐标; 单元定义时类自动将y,z坐标设为0;
    
    载荷支持集中节点力,不考虑热效应
    """
    VTK_TYPE=3
    def __init__(self,label:int,
                    node1:Node,node2:Node,
                    E=0.0,A=0.0,rho=0.0):
        super().__init__(label,1)
        self.node1=node1
        self.node2=node2
        self.E=E
        self.A=A
        self.rho=rho
        # y，z坐标强制设为0
        self.node1.y,self.node1.z=0.0,0.0
        self.node2.y,self.node2.z=0.0,0.0
    def __str__(self):
        return f"T1D2({self.label},{self.node1.label},{self.node2.label}),E={self.E},A={self.A},rho={self.rho}"
    @property
    def get_connectivity(self):
        return [self.node1.label,self.node2.label]
    @property
    def L(self):
        """
        计算单元长度
        
        returns:
            float: 单元长度
        """
        return abs(self.node2.x-self.node1.x)
    @property
    def Ke(self):
        """
        计算当前单元的全局坐标系下的单元刚度矩阵
        
        returns:
            np.ndarray: 单元刚度矩阵(2,2)
        """
        return np.array([[1,-1],[-1,1]])*(self.E*self.A/self.L)
    @property
    def Me(self):
        """
        计算当前单元的集中质量矩阵(总质量均布到每个节点)
        
        returns:
            np.ndarray: 单元质量矩阵(2,2)
        """
        return (self.rho*self.A*self.L/2.0)*np.eye(2)
    @property
    def D(self):
        """
        返回当前单元的弹性模量E
        
        returns:
            float: 弹性模量E
        """
        return self.E
    @property
    def dofNum(self):
        """
        获取当前单元的自由度数量
        
        returns:
            int: 自由度数量
        """
        return 2
    @property
    def NodesList(self):
        """
        获取当前单元的节点列表
        
        returns:
            List[Node]: 节点列表
        """
        return [self.node1,self.node2]
    @property
    def ElemDofIndexs(self):
        """
        获取当前单元的全局自由度索引列表(0-based)
        
        returns:
            List[int]: 自由度索引列表
        """
        return dofInds(self.node1.label,self.NodeDofNum)+dofInds(self.node2.label,self.NodeDofNum)
    def AxialStress(self,ue=None):
        """
        计算单元的轴向应力
        
        args:
            ue (np.ndarray, optional): 单元位移向量(2,). Defaults to None.
        
        returns:
            float: 轴向应力
        """
        if ue is None:
            raise ValueError("ue must be provided to calculate axial stress")
        return self.AxialStrain(ue)*self.E
    def AxialStrain(self,ue=None):
        """
        计算单元的轴向应变
        Args:
            ue (np.ndarray, optional): 单元位移向量(2,). Defaults to None.
        
        returns:
            float: 轴向应变
        """
        if ue is None:
            raise ValueError("ue must be provided to calculate axial strain")
        return (ue[1]-ue[0])/self.L